#!/usr/bin/env python

import csv
import cv2
import numpy as np
from cv_bridge import CvBridge, CvBridgeError

import rospy
from sensor_msgs.msg import Image, JointState
from interbotix_xs_msgs.msg import JointGroupCommand
from interbotix_xs_msgs.srv import RobotInfo

### Class that performs object tracking based on color segmentation
class ColorTracker(object):

    ### @brief Initialization of the ColorTracker class; sets up the components necessary to do object tracking
    def __init__(self):
        hsv_config_path = rospy.get_param('~hsv_configs')                                                               # ROS Paramter containing the path to the 'hsv.csv' file with the HSV parameters
        with open(hsv_config_path) as csv_file:                                                                         # Open the CSV file
            csv_reader = list(csv.reader(csv_file))                                                                     # Load the file contents into the program
            self.lower_threshold = np.array([int(csv_reader[0][1]), int(csv_reader[2][1]), int(csv_reader[4][1])])      # Set the lower threshold for HSV color segmentation
            self.upper_threshold = np.array([int(csv_reader[1][1]), int(csv_reader[3][1]), int(csv_reader[5][1])])      # Set the upper threshold for HSV color segmentation
        self.center_x = 320/2.0                                                                                         # Middle pixel 'x' position
        self.center_y = 240/2.0                                                                                         # Middle pixel 'y' position
        self.cv_image = None                                                                                            # Raw image in a format that OpenCV understands
        self.joint_states = None                                                                                        # Current states of the 'pan' and 'tilt' joints
        self.bridge = CvBridge()                                                                                        # Converts between Sensor_msgs/Image and OpenCV image formats
        rospy.wait_for_service("get_robot_info")                                                                        # Wait for the 'get_robot_info' ROS Servie to become available
        self.srv_robot_info = rospy.ServiceProxy("get_robot_info", RobotInfo)                                           # ROS Service used to get joint limit information
        self.pub_cmds = rospy.Publisher("commands/joint_group", JointGroupCommand, queue_size=1)                        # ROS Publisher to publish commands to the 'pan' and 'tilt' joints
        self.publish_contour_detections_image = rospy.get_param('~publish_contour_detections_image')                    # ROS Parameter containing a boolean on whether the node should publish images showing the detected contours
        if (self.publish_contour_detections_image):
            self.pub_images = rospy.Publisher("contour_detections", Image, queue_size=5)                                # ROS Publisher to publish images with the detected contours
        self.sub_joint_states = rospy.Subscriber("joint_states", JointState, self.joint_state_cb)                       # ROS Subscriber to get the latest joint states
        while (self.joint_states is None): pass                                                                         # Wait until the node gets the latest joint states
        self.joint_commands = JointGroupCommand("turret", list(self.joint_states.position))                             # Set initial joint commands to the current joint positions
        self.robot_info = self.srv_robot_info("group", "turret")                                                        # Call the RobotInfo Service
        self.sub_img = rospy.Subscriber("lifecam/image_raw", Image, self.image_cb)                                      # ROS Subscriber to get images from the 'usb_cam' node

    ### @brief ROS Subscriber callback function to get updated joint states
    ### @param msg - updated JointState message
    def joint_state_cb(self, msg):
        self.joint_states = msg

    ### @brief ROS Subscriber callback function to get the latest images from the 'usb_cam' node
    ### @param msg - updated Image message
    def image_cb(self, msg):
        try:
          cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except CvBridgeError as e:
          print(e)

        # convert to the HSV colorspace
        hsv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2HSV)
        # create the mask
        cv_mask = cv2.inRange(hsv_image, self.lower_threshold, self.upper_threshold)
        # use a Morphological Operation to clean up 'noise' as explained at https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_morphological_ops/py_morphological_ops.html#opening
        # An Ellipsoid kernel seems to do the job
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
        cv_mask = cv2.morphologyEx(cv_mask, cv2.MORPH_OPEN, kernel)
        # create the overlay image (show the entire image overlaid with 'yellow' points that define the boundary of interest and a 'red' point that defines the centroid within that boundary)
        cv_overlay = hsv_image.copy()
        contours, _ = cv2.findContours(cv_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
        # if a 'contour' (fancy word for boundary) was found...
        if len(contours) > 0:
            # return the largest contour from all the ones found
            c = max(contours, key = cv2.contourArea)
            # find the approximate center of that contour
            ((x,y), radius) = cv2.minEnclosingCircle(c)
            # draw a red circle at the center
            cv2.circle(cv_overlay,(int(x),int(y)),5,(0,255,255),-1)
            # draw yellow points to mark the boundary
            cv2.drawContours(cv_overlay,c,-1, (30,255,255),3)
            # Publish JointCommands so that the center pixel in the image overlaps the red circle mentioned above
            if ((self.center_x - x) > 10):
                self.joint_commands.cmd[0] += 0.02
                if (self.joint_commands.cmd[0] > self.robot_info.joint_upper_limits[0]):
                    self.joint_commands.cmd[0] = self.robot_info.joint_upper_limits[0]
            elif ((self.center_x - x) < -10):
                self.joint_commands.cmd[0] -= 0.02
                if (self.joint_commands.cmd[0] < self.robot_info.joint_lower_limits[0]):
                    self.joint_commands.cmd[0] = self.robot_info.joint_lower_limits[0]
            if ((self.center_y - y) > 10):
                self.joint_commands.cmd[1] -= 0.02
                if (self.joint_commands.cmd[1] < self.robot_info.joint_lower_limits[1]):
                    self.joint_commands.cmd[1] = self.robot_info.joint_lower_limits[1]
            elif ((self.center_y - y) < -10):
                self.joint_commands.cmd[1] += 0.02
                if (self.joint_commands.cmd[1] > self.robot_info.joint_upper_limits[1]):
                    self.joint_commands.cmd[1] = self.robot_info.joint_upper_limits[1]
            self.pub_cmds.publish(self.joint_commands)

        # if the overlay image should be published...
        if (self.publish_contour_detections_image):
            cv_overlay = cv2.cvtColor(cv_overlay, cv2.COLOR_HSV2BGR)
            try:
                self.pub_images.publish(self.bridge.cv2_to_imgmsg(cv_overlay, "bgr8"))
            except CvBridgeError as e:
                print(e)

if __name__ == '__main__':
    rospy.init_node('xsturret_color_tracker')
    color_tracker = ColorTracker()
    rospy.spin()
